import faiss
import numpy #科学计算库，它提供了高性能的多维数组对象以及对这些数组进行操作的工具。

#基本操作
def fun_1():
    #创建索引
    index = faiss.IndexFlatL2(768) # 向量的维度 ， flat表示线性搜索 L2表示相似度的计算是欧式距离，越小表示越相似
    index = faiss.IndexFlatIP(768) # 向量的维度 ， flat表示线性搜索 IP表示相似度的计算是内积，越大表示越相似
    index = faiss.index_factory(768, "Flat" , faiss.METRIC_L2) # 使用工厂方法创建 ， 等同于IndexFlatL2
    index = faiss.index_factory(768, "Flat" , faiss.METRIC_INNER_PRODUCT) # 使用工厂方法创建 ， 等同于IndexFlatIP


    #添加向量
    vectors = numpy.random.random(size=(1000, 768)) # 创建1000个随机向量 , 768表示向量的维度
    index.add(vectors) #注意，这种方式添加向量，会自动给每个向量一个索引，从0开始


    #搜索向量
    query = numpy.random.random(size=(2, 768)) # 创建2个随机向量
    D , I = index.search(query, 2) # 搜索前5个最相似的向量 , D表示相似度，I表示索引
    print(D)
    print(I)


    #删除向量
    index.remove_ids(numpy.array([1, 2])) # 删除索引为1和2的向量
    print(index.ntotal) # 索引中向量的数量

    index.reset() # 删除所有向量
    print(index.ntotal)  # 索引中向量的数量

    #存储索引
    faiss.write_index(index, "index.faiss") # 保存索引到文件


    #加载索引
    index = faiss.read_index("index.faiss") #从文件加载索引
    print(index)


#ID映射
def fun_2():
    # 使用 IndexIDMap2 来支持 add_with_ids 操作
    index = faiss.IndexFlatIP(768)
    index = faiss.IndexIDMap(index)  # 使用 IndexIDMap让索引支持手动向量ID

    vectors = numpy.random.random(size=(1000, 768)).astype('float32')  # 转换为 float32 以符合 Faiss 要求
    ids = numpy.arange(1000, 2000).astype('int64')  # Faiss 要求 ID 为 int64 类型
    index.add_with_ids(vectors, ids)  # 向量ID从1000开始
    print(index.ntotal)

    index.remove_ids(numpy.array([0,1,2,3])) # 删除索引为1,2,3的向量
    print(index.ntotal) #发现向量没有减少

    #有些索引类型本身支持用户指定ID，如果不支持可以使用IndexIDMap包装，以支持用户指定ID
if __name__ == '__main__':
    #fun_1()
    fun_2()